package io.kiki.sba.registry.server.clientInterface.registry;


import io.kiki.sba.registry.api.Channel;
import io.kiki.sba.registry.api.exchange.Callback;
import io.kiki.sba.registry.common.model.dataserver.DatumVersion;
import io.kiki.sba.registry.server.clientInterface.bootstrap.Config;
import io.kiki.sba.registry.server.clientInterface.node.service.DataStoreNodeService;
import io.kiki.sba.registry.server.clientInterface.slot.SlotsCache;
import io.kiki.sba.registry.util.ConcurrentUtils;
import lombok.Getter;
import lombok.Setter;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.Map.Entry;
import java.util.concurrent.TimeUnit;


public class RegistryScanCallable {

    private static final Logger logger = LoggerFactory.getLogger(RegistryScanCallable.class);

    @Autowired
    private Config config;

    @Autowired
    private SlotsCache slotsCache;

    /**
     * transfer data to DataNode
     */
    @Autowired
    private DataStoreNodeService dataStoreNodeService;

    public void scanVersions(long round, String dataCenter, Map<String, DatumVersion> dataInfoIdToVersioMap, ScanCallable scanCallable) {
        Map<Integer, Map<String, DatumVersion>> slotIdToDataInfoIdToVersioMap = groupBySlot(dataInfoIdToVersioMap);

        Map<Integer, FetchVersionResult> slotIdToFetchVersionResultMap = new HashMap<>(dataInfoIdToVersioMap.size());
        for (Map.Entry<Integer, Map<String, DatumVersion>> entry : slotIdToDataInfoIdToVersioMap.entrySet()) {
            final Integer slotId = entry.getKey();
            try {
                final FetchVersionResult fetchVersionResult = fetchDataVersionAsync(dataCenter, slotId, entry.getValue(), round);
                if (fetchVersionResult != null) {
                    slotIdToFetchVersionResultMap.put(slotId, fetchVersionResult);
                }
            } catch (Throwable e) {
                logger.info("[fetchSlotVer]round={},{},{},leader={},interests={},gets={},success={}", round, slotId, dataCenter, slotsCache.getLeader(dataCenter, slotId), dataInfoIdToVersioMap.size(), 0, "N");

                logger.error("round={}, failed to fetch versions slotId={}, size={}", round, slotId, entry.getValue().size(), e);
            }
        }
        final int timeoutMillis = config.getDataNodeExchangeTimeoutMillis();
        final long waitDeadline = System.currentTimeMillis() + timeoutMillis + 2000;
        // wait async finish
        ConcurrentUtils.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
        // check callback result, use for.count to avoid the clock skew
        for (int i = 0; i < (timeoutMillis * 2) / 50; i++) {
            handleFetchResult(round, dataCenter, slotIdToFetchVersionResultMap, scanCallable);
            if (slotIdToFetchVersionResultMap.isEmpty() || System.currentTimeMillis() > waitDeadline) {
                break;
            }
            ConcurrentUtils.sleepUninterruptibly(50, TimeUnit.MILLISECONDS);
        }
        handleFetchResult(round, dataCenter, slotIdToFetchVersionResultMap, scanCallable);
        if (!slotIdToFetchVersionResultMap.isEmpty()) {
            logger.error("[fetchSlotVerTimeout]round={},callbacks={},{}", round, slotIdToFetchVersionResultMap.size(), slotIdToFetchVersionResultMap);
        }
    }

    int handleFetchResult(long round, String dataCenter, Map<Integer, FetchVersionResult> slotIdToFetchVersionResultMap, ScanCallable scanCallable) {
        int count = 0;
        final Iterator<Entry<Integer, FetchVersionResult>> iterator = slotIdToFetchVersionResultMap.entrySet().iterator();
        while (iterator.hasNext()) {
            final Map.Entry<Integer, FetchVersionResult> entry = iterator.next();
            FetchVersionResult fetchVersionResult = entry.getValue();
            if (fetchVersionResult.fetchVersionCallback == null) {
                // not finish
                continue;
            }
            iterator.remove();
            count++;
            // success
            if (fetchVersionResult.fetchVersionCallback.dataInfoIdToDatumVersionMap != null) {
                final long now = System.currentTimeMillis();

                // execute
                fetchVersionResult.fetchVersionCallback.dataInfoIdToDatumVersionMap.forEach((dataInfoId, datumVersion) -> scanCallable.execute(new ScanCallableInfo(round, dataCenter, entry.getKey(), fetchVersionResult.leader, dataInfoId, datumVersion, now)));
            }
        }
        return count;
    }

    private Map<Integer, Map<String, DatumVersion>> groupBySlot(Map<String, DatumVersion> dataInfoIdToVersioMap) {
        if (CollectionUtils.isEmpty(dataInfoIdToVersioMap)) {
            return Collections.emptyMap();
        }
        TreeMap<Integer, Map<String, DatumVersion>> map = new TreeMap<>();
        for (Map.Entry<String, DatumVersion> entry : dataInfoIdToVersioMap.entrySet()) {
            final String dataInfoId = entry.getKey();
            Map<String, DatumVersion> dataInfoIdToVersioMap2 = map.computeIfAbsent(slotsCache.slotOf(dataInfoId), k -> new HashMap<>(256));
            dataInfoIdToVersioMap2.put(dataInfoId, entry.getValue());
        }
        return map;
    }

    FetchVersionResult fetchDataVersionAsync(String dataCenter, int slotId, Map<String, DatumVersion> interestVersions, long round) {
        final String leader = slotsCache.getLeader(dataCenter, slotId);
        if (StringUtils.isBlank(leader)) {
            logger.error("[NoLeader]slotId={}, round={}", slotId, round);
            return null;
        }
        final FetchVersionResult fetchVersionResult = new FetchVersionResult(slotId, leader);
        dataStoreNodeService.fetchDataVersion(dataCenter, slotId, interestVersions, new Callback<Map<String, DatumVersion>>() {
            @Override
            public void onSuccess(Channel channel, Map<String, DatumVersion> dataInfoIdToDatumVersionMap) {
                // merge the version
                Map<String, DatumVersion> map = new HashMap<>(interestVersions);
                map.putAll(dataInfoIdToDatumVersionMap);
                fetchVersionResult.fetchVersionCallback = new FetchVersionCallback(map);
                logger.info("[fetchSlotVer]round={},{},{},leader={},interests={},gets={},success={}", round, slotId, dataCenter, leader, interestVersions.size(), dataInfoIdToDatumVersionMap.size(), "Y");
            }

            @Override
            public void onFailed(Channel channel, Throwable e) {
                fetchVersionResult.fetchVersionCallback = new FetchVersionCallback(null);
                logger.info("[fetchSlotVer]round={},{},{},leader={},interests={},gets={},success={}", round, slotId, dataCenter, leader, interestVersions.size(), 0, "N");

                logger.error("round={},failed to fetch versions,slotId={},leader={},size={}", round, slotId, leader, interestVersions.size(), e);
            }
        });
        return fetchVersionResult;
    }

    public interface ScanCallable {
        void execute(ScanCallableInfo callableInfo);
    }

    private static final class FetchVersionResult {
        final String leader;
        final int slotId;
        volatile FetchVersionCallback fetchVersionCallback;

        FetchVersionResult(int slotId, String leader) {
            this.leader = leader;
            this.slotId = slotId;
        }


    }

    private static final class FetchVersionCallback {
        final Map<String, DatumVersion> dataInfoIdToDatumVersionMap;

        FetchVersionCallback(Map<String, DatumVersion> dataInfoIdToDatumVersionMap) {
            this.dataInfoIdToDatumVersionMap = dataInfoIdToDatumVersionMap;
        }
    }

    @Setter
    @Getter
    public final class ScanCallableInfo {
        private final long round;
        private final String dataCenter;
        private final int slotId;
        private final String leader;

        private final String dataInfoId;

        private final DatumVersion datumVersion;
        private final long currentTs;

        public ScanCallableInfo(long round, String dataCenter, int slotId, String leader, String dataInfoId, DatumVersion datumVersion, long currentTs) {
            this.round = round;
            this.dataCenter = dataCenter;
            this.slotId = slotId;
            this.leader = leader;
            this.dataInfoId = dataInfoId;
            this.datumVersion = datumVersion;
            this.currentTs = currentTs;
        }


    }
}
